import numpy as np
import torch
import tensorflow as tf
from tensorflow import keras
# import keras
# 此处有些版本的tensorflow用第一种方式导入keras，有些用第二种方式
import onnx
import onnxruntime

np.random.seed(0)
tf.random.set_seed(0)
torch.random.manual_seed(0)


def build_tf_dense(w: np.ndarray, b: np.ndarray) -> keras.layers.Layer:
    # 使用w和b构造tensorflow的全连接层
    M, N = w.shape
    layer = keras.layers.Dense(N)
    layer.build((1, 1, M))
    layer.weights[0].assign(w)
    layer.weights[1].assign(b)
    return layer


def build_torch_linear(w: np.ndarray, b: np.ndarray) -> torch.nn.Module:
    # 使用w和b构造torch的全连接层
    M, N = w.shape
    layer = torch.nn.Linear(M, N)
    layer.weight.data = torch.from_numpy(w.T)
    layer.bias.data = torch.from_numpy(b)
    return layer


def build_onnx_fc(w: np.ndarray, b: np.ndarray):
    # 使用w和b构造onnx的全连接层
    matmul_node = onnx.helper.make_node('MatMul', ['x', 'w'], ['xw'], name='matmul_fc')
    add_node = onnx.helper.make_node('Add', ['xw', 'b'], ['y'], name='add_fc')
    w_init = onnx.helper.make_tensor('w', onnx.TensorProto.FLOAT, w.shape, w)
    b_init = onnx.helper.make_tensor('b', onnx.TensorProto.FLOAT, b.shape, b)
    x = onnx.helper.make_tensor_value_info('x', onnx.TensorProto.FLOAT, None)
    y = onnx.helper.make_tensor_value_info('y', onnx.TensorProto.FLOAT, None)
    graph = onnx.helper.make_graph([matmul_node, add_node], 'fullconnect',
                                   [x], [y], [w_init, b_init])
    model = onnx.helper.make_model(graph, opset_imports=[onnx.helper.make_opsetid("", 15)])
    model.ir_version = 6

    # onnx.checker.check_model(model)
    # 注意此处check是通过不了的，因为没有设置输入输出的维度，不设置维度是为了适应多种输入尺寸
    onnx.save(model, 'onnx_hub/full_connect.onnx')


def onnx_infer_fc(x: np.ndarray) -> np.ndarray:
    # 使用onnxrumtime推理onnx文件
    onnxmodel = onnxruntime.InferenceSession('onnx_hub/full_connect.onnx')
    y = onnxmodel.run(['y'], {'x': x})[0]
    return y


if __name__ == '__main__':
    M = 5
    N = 4
    weight = np.random.random((M, N)).astype('float32')  # 构造一个全连接的权重和偏置
    bias = np.random.random((N,)).astype('float32')

    layer_tf = build_tf_dense(weight, bias)
    layer_torch = build_torch_linear(weight, bias)
    build_onnx_fc(weight, bias)

    x = np.random.random((1, 2, 3, 7, M)).astype('float32')  # M前面的维度可以任意设置，对全连接都是合法的
    y0 = np.matmul(x, weight) + bias  # numpy实现的全连接
    y1 = layer_tf(x).numpy()
    y2 = layer_torch(torch.from_numpy(x)).detach().numpy()
    y3 = onnx_infer_fc(x)

    # 评估三种形式的误差情况
    err1 = np.mean((y1 - y0) ** 2)
    err2 = np.mean((y2 - y0) ** 2)
    err3 = np.mean((y3 - y0) ** 2)
    print(f'err1:{err1}, err2:{err2}, err3:{err3}')
    pass
